# 锁定跟踪之CamShift算法
# 原理：利用目标的颜色直方图进行反向投影，通过meanshift(均值漂移)迭代找到目标位置

import cv2
import numpy

# 读取视频
cap = cv2.VideoCapture(0)

# 读取第一帧
ret, frame = cap.read()
if not ret:
    print("read error\n")
    exit(0)

cv2.imshow("frame", frame)

width, height, channels = frame.shape

# 设置初始窗口
x, y, w, h = cv2.selectROI("frame", frame)
track_window = (x, y, w, h)
print(track_window)

# 设置 ROI (Region of Interest)
roi = frame[y:y+h, x:x+w]
cv2.imshow("roi", roi)

# 转换为 HSV 颜色空间
hsv_roi = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
cv2.imshow("hsv_roi", hsv_roi)

# hsv上下界（一维数组）
lower = numpy.array((0.0, 60.0, 32.0))
upper = numpy.array((180.0, 255.0, 255.0))

# 创建掩膜
# 判断hsv_roi是否在给定的range里，结果返回到mask(大小与src相同，二值图像)，满足range的值为255，否则为0
mask = cv2.inRange(hsv_roi, lower, upper)
cv2.imshow("mask", mask)

# 计算直方图
roi_hist = cv2.calcHist([hsv_roi], [0], mask, [180], [0, 180])

cv2.normalize(roi_hist, roi_hist, 0, 255, cv2.NORM_MINMAX)
cv2.imshow("roi_hist", roi_hist)

# 设置终止条件
term_crit = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 1)

while True:
    ret, frame = cap.read()
    if not ret:
        break

    # 转换为 HSV 颜色空间
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

    # 计算反向投影
    dst = cv2.calcBackProject([hsv], [0], roi_hist, [0, 180], 1)

    # 应用 CamShift 算法
    ret, track_window = cv2.CamShift(dst, track_window, term_crit)

    # 绘制跟踪结果
    pts = cv2.boxPoints(ret)
    pts = numpy.int64(pts)
    img2 = cv2.polylines(frame, [pts], True, 255, 2)
    cv2.imshow('CamShift Tracking', img2)

    if cv2.waitKey(30) & 0xFF == 27:
        break

cap.release()
cv2.destroyAllWindows()